#!/usr/bin/env python3
# pur_policy_client.py
# 远程策略客户端：连接pur_policy服务器，加载DRL模型，接收状态并返回计算结果
import socket
import struct
import numpy as np
import sys
import logging
from drl_policy_driver import DRLPolicy

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def receive_all(sock, length):
    """可靠地接收指定长度的数据"""
    data = b''
    while len(data) < length:
        packet = sock.recv(length - len(data))
        if not packet:
            raise ConnectionError("连接中断")
        data += packet
    return data

# 在客户端使用这个函数接收邻接矩阵数据
# adj_bytes = receive_all(client_socket, adj_length)

def main():
    server_ip = "127.0.0.1"
    model_path = "./pursuer/model/pursuer_model/policy.pth"
    port = 1217

    try:
        # 连接到策略服务器
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.connect((server_ip, port))
        logger.info(f"已连接到策略服务器 {server_ip}:{port}")

        # 接收节点数
        node_num_data = receive_all(sock, 4)
        node_num = struct.unpack(">I", node_num_data)[0]
        logger.info(f"接收到节点数: {node_num}")

        # 接收邻接矩阵长度
        len_data = receive_all(sock, 4)
        len_adj = struct.unpack(">I", len_data)[0]
        logger.info(f"接收到邻接矩阵长度: {len_adj}")

        # 接收邻接矩阵字节
        adj_bytes = receive_all(sock, len_adj)
        logger.info(f"接收到邻接矩阵字节,大小为 {len(adj_bytes)}")
        adj_matrix = np.frombuffer(adj_bytes, dtype=np.uint8).reshape(
            (node_num, node_num)
        )
        logger.info("接收到邻接矩阵")

        # 加载DRL模型
        drl_policy = DRLPolicy(adj_matrix, model_path)
        logger.info(f"DRL模型已加载: {model_path}")

        # 主循环：接收请求并响应
        while True:
            # 接收命令
            cmd = sock.recv(1)
            if not cmd:
                break
            if cmd == b"\x01":  # 请求计算下一步
                # 接收当前节点状态
                state_data = sock.recv(12)
                if len(state_data) != 12:
                    logger.warning("状态数据不完整，跳过")
                    continue
                evader_node, p1_node, p2_node = struct.unpack(">III", state_data)
                logger.debug(
                    f"接收到状态: 逃避者={evader_node}, 追捕者1={p1_node}, 追捕者2={p2_node}"
                )

                # 调用DRL计算
                next_evader = evader_node
                next_p1, next_p2 = drl_policy.calculate_next_move(
                    evader_node, p1_node, p2_node
                )

                # 发送响应
                response = struct.pack(">III", next_evader, next_p1, next_p2)
                sock.sendall(response)
                logger.debug(
                    f"发送策略: 逃避者={next_evader}, 追捕者1={next_p1}, 追捕者2={next_p2}"
                )
            else:
                logger.warning(f"未知命令: {cmd}")

    finally:
        if "sock" in locals():
            sock.close()
        logger.info("客户端已关闭")


if __name__ == "__main__":
    main()
